"""
 * @file main.py
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-11-3
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi. This file can not
 * be copied or distributed without the permission of Lynxi Technologies Co., Ltd.
"""
import pylynchipsdk as sdk
import numpy as np
import getopt

import sys
sys.path.append('../../../')
import common.python.common as common
import common.python.bufferpool as bufferpool

class decoder_attr:
    def __init__(self) -> None:
        self.target_path = ""
        self.source_path = ""
        self.params_path = ""
        self.model_path = ""
        self.device_id = 0

def usage():
    print("usage: python3 main.py")
    print("\t -t    target path, eg: ../../../data/multiple_input_output_model_infer_sample/All_Clean_2023-06-13_10_32.jpg")
    print("\t -s    source path, eg: ../../../data/multiple_input_output_model_infer_sample/All_Clean_2023-06-13_12_01_27.789072.jpg")
    print("\t -p    params path, eg: ../../../data/multiple_input_output_model_infer_sample/params-A-new.json")
    print("\t -m    model path, eg: ../../../model/pspnet_hrnet_w18_input_128_float16_11.80_/Net_0")
    print("\t -d    device id, default 0")
    print("\t -h    show help")
    print("\t -?    show help")

def main():

    ## 输入 
    input_data_1 = common.imgRead(attr.target_path)
    input_data_2 = common.imgRead(attr.source_path)
    params = common.jsonLoad(attr.params_path)['11-advertisement']
    rotations, boxes = params['rotations'], params['boxes']
    ## 创建context
    context, ret = sdk.lyn_create_context(attr.device_id)
    common.error_check(ret, "lyn_create_context")
    ret = sdk.lyn_register_error_handler(common.default_stream_error_handle, None)
    common.error_check(ret, "lyn_register_error_handler")
    # 创建 stream
    apu_stream, ret = sdk.lyn_create_stream()
    common.error_check(ret, "lyn_create_stream")

    # 加载模型
    model,ret = sdk.lyn_load_model(attr.model_path)
    common.error_check(ret, "lyn_load_model")
    model_desc, ret = sdk.lyn_model_get_desc(model)
    common.error_check(ret, "lyn_model_get_desc")
    batch_size,channels,model_width,model_height = model_desc.inputTensorAttrArray[0].dims
    ipe_output_size = model_width * model_height * batch_size * channels
    model_output_size = model_desc.outputDataLen
    model_intput_size = model_desc.inputDataLen
    model_output_nums = model_desc.outputTensorAttrArrayNum
    # print("model_output_nums:",model_output_nums)

    # 数据前处理
    input_data1_list = common.preprocess(input_data_1,rotations,boxes,(128,128))
    input_data2_list = common.preprocess(input_data_2,rotations,boxes,(128,128))

    # 为模型输入申请device内存
    model_input, ret = sdk.lyn_malloc(model_intput_size)
    common.error_check(ret, "lyn_malloc")
    input_data1_ptr = sdk.lyn_numpy_to_ptr(input_data1_list[0])
    input_data2_ptr = sdk.lyn_numpy_to_ptr(input_data2_list[0])
    ret = sdk.lyn_memcpy(model_input, input_data1_ptr, input_data1_list[0].nbytes, sdk.lyn_memcpy_dir_t.ClientToServer)
    common.error_check(ret, "lyn_memcpy")
    new_addr = sdk.lyn_addr_seek(model_input,input_data1_list[0].nbytes)
    ret = sdk.lyn_memcpy(new_addr, input_data2_ptr, input_data2_list[0].nbytes, sdk.lyn_memcpy_dir_t.ClientToServer)
    common.error_check(ret, "lyn_memcpy")
    apu_output_mem_pool = bufferpool.buffer_pool(model_output_size * batch_size, len(input_data1_list))
    # # apu 推理
    apu_output = apu_output_mem_pool.pop()
    ret = sdk.lyn_execute_model_async(apu_stream, model, model_input, apu_output, batch_size)
    common.error_check(ret, "lyn_execute_model_async")

    # 等待流处理完
    ret = sdk.lyn_synchronize_stream(apu_stream)
    common.error_check(ret, "lyn_synchronize_stream")
    
    result = []
    offset = 0
    off_address = apu_output

    for i in range(model_output_nums):
        shape, ret = sdk.lyn_model_get_output_tensor_dims_by_index(model, i)
        common.error_check(ret, "lyn_model_get_input_tensor_dims_by_index")
        dtype = model_desc.outputTensorAttrArray[i].dtype
        size = model_desc.outputTensorAttrArray[i].dataLen

        out_data = np.empty(int(size/2), dtype=np.float16)
        out_data_ptr = sdk.lyn_numpy_to_ptr(out_data) 

        sdk.lyn_memcpy(out_data_ptr, off_address, size, sdk.lyn_memcpy_dir_t.ServerToClient)
        result.append(out_data)

        offset = offset + size
        off_address = sdk.lyn_addr_seek(apu_output, offset) 

    ss_0 = np.reshape(result[0],[-1, 6, 128, 128])
    ss_1 = np.reshape(result[1],[-1, 6, 128, 128])
    ss_2 = np.reshape(result[2],[-1, 128, 128])
    #打印三个输出
    print("result---------0----",ss_0)
    print("result---------1----",ss_1)
    print("result---------2----",ss_2)

    # np.save("./lynxi_output_0.npy",ss_0)
    # np.save("./lynxi_output_1.npy",ss_1)
    # np.save("./lynxi_output_2.npy",ss_2)

    # 销毁stream
    ret = sdk.lyn_destroy_stream(apu_stream)
    common.error_check(ret, "lyn_destroy_stream")
    # 释放内存
    ret = sdk.lyn_free(model_input)
    common.error_check(ret, "lyn_free")
    ## 卸载模型
    ret = sdk.lyn_unload_model(model)
    common.error_check(ret, "lyn_unload_model")
    # 销毁 context
    ret = sdk.lyn_destroy_context(context)
    common.error_check(ret, "lyn_destroy_context")

if __name__ == "__main__":

    argv = sys.argv[1:]
    try:
        opts, args = getopt.getopt(argv, "t:s:p:m:d:h?")
    except getopt.GetoptError:
        usage()
        sys.exit(0)
    
    attr = decoder_attr()

    for opt, arg in opts:
        if opt in("-h", "-?"):
            usage()
            sys.exit(0)
        elif opt == "-t":
            attr.target_path = arg
        elif opt == "-s":
            attr.source_path = arg
        elif opt == "-p":
            attr.params_path = arg
        elif opt == "-m":
            attr.model_path = arg
        elif opt == "-d":
            attr.device_id = int(arg)
        else:
            usage()
            sys.exit(0)

    if (attr.target_path == ""):
        print("Plase input target path")
        usage()
        sys.exit(0)
    elif (attr.source_path == ""):
        print("Plase input source path")
        usage()
        sys.exit(0)
    elif (attr.params_path == ""):
        print("Plase input params path")
        usage()
        sys.exit(0)
    elif (attr.model_path == ""):
        print("Plase input model path")
        usage()
        sys.exit(0)

    main()
